from typing import Callable, List, Iterator, Optional
from pprint import pprint
from numpy.lib.utils import source
from random import shuffle

import pandas as pd
import numpy as np
import heapq
import os
import sys


def read_int_data(filename: str) -> np.ndarray:
    data = pd.read_csv(filename, header=None)
    return data.to_numpy()


def euclidean_distance(vector1: np.ndarray, vector2: np.ndarray) -> float:
    """
    求两个向量的欧式距离
    """
    return np.sqrt(np.sum((vector1 - vector2)**2))

class Index:
    def __init__(self, index: int, dist_to_source: float) -> None:
        self.index = index
        self.dist_to_source = dist_to_source

    def __lt__(self, ano: "Index") -> bool:
        """
        比较vectors中下标为self.index的点 与下标为ano.index的点分别与source点的距离
        """
        if self.dist_to_source != ano.dist_to_source:
            return self.dist_to_source < ano.dist_to_source
        return self.index < ano.index

    def __repr__(self) -> str:
        return "[{:d}] {:.4f}".format(self.index, self.dist_to_source)

def k_nearest_of_point(source: np.ndarray,
                       vectors: List[np.ndarray],
                       distance: Callable[[np.ndarray, np.ndarray], float],
                       k: Optional[int] = None) -> List[int]:
    """
    在vectors中找到source的k近邻点

    source: np.ndarray，源点坐标
    vectors: List[np.ndarray]，源点source的邻近点集合
    distance: 可调用的函数对象，用于计算source和vectors[i]的距离
    k: int，在vectors中选取至多k个、且距离source最近的点

    returns: 距离源点source的k个最近点在vectors中的下标  
    """

    if k is None or k < 0:
        k = len(vectors)
    index_wrappers = [Index(i, distance(source, vectors[i]))
                      for i in range(len(vectors))]
    res_index = heapq.nsmallest(k, index_wrappers)
    res = [obj.index for obj in res_index]  # 取出下标
    return res


def k_nearest_matrix(matrix: np.ndarray,
                     distance: Callable[[np.ndarray, np.ndarray], float],
                     k: Optional[int] = None) -> List[List[int]]:
    """
    求matrix中每个点在matrix中的k近邻点

    matrix: np.ndarray，二维矩阵，每一行表示一个点的坐标
    distance: 可调用对象，可用于求matrix中任意两点的距离
    k: int，求至多k个近邻

    returns: List[List[int]]，res[i]为List[int]，表示第matrix[i]点的所有近邻点在matrix中的下标，且按照距离递增顺序排序，res[i][0]是点本身的下标
    """
    res = []  # type: List[List[int]]
    point_list = list(matrix)  # type: List[np.ndarray]
    res = []
    for source in point_list:
        indeces = k_nearest_of_point(source, point_list,
                                     distance, k)  # type: List[int]
        res.append(indeces)
    return res


def get_step_increment_of_farthest_neighbor(matrix: np.ndarray) -> List[List[float]]:
    """
    求matrix中每个点的最远近邻点的增长步长，即，k=i+1近邻点和源点的距离 - k=i近邻点和源点的距离，其中i从1开始                      
    """
    neighbors_indeces = k_nearest_matrix(
        matrix, distance=euclidean_distance)  # type: List[List[int]]
    res = []  # type: List[List[float]]
    for indeces in neighbors_indeces:
        current_point_index = indeces[0]
        current_point = matrix[current_point_index]
        distance_list = []  # type: List[float]
        for neighbor_index in indeces:
            neighbor = matrix[neighbor_index]  # type: np.ndarray
            distance = euclidean_distance(
                current_point, neighbor)  # 求出当前点和邻居的距离
            distance_list.append(distance)
        steps = []  # List[float]
        for i in range(2, len(distance_list)):
            # k=i+1近邻点和源点的距离-k=i近邻点和源点的距离，定义为近邻点距离的增加步长
            step = distance_list[i] - distance_list[i - 1]
            steps.append(step)
        res.append(steps)

    return res

def vector_module(vec: np.ndarray) -> int:
    """
    求向量的模
    """
    ndimension = len(vec)
    return euclidean_distance(vec, np.zeros(shape=(ndimension)))

def get_step_increment_of_gravity(matrix: np.ndarray) -> List[List[float]]:
    """
    求每个点随着k取值的增大，近邻点对源点产生合力大小的增长步长
    """
    res = [] # List[List[float]]
    neighbors_indeces = k_nearest_matrix(
        matrix, distance=euclidean_distance)  # type: List[List[int]]
    for i, source in enumerate(matrix):
        neighbor_info = neighbors_indeces[i] # 当前点的临近点信息
        gravity_vector_size_list = [] # 每次增加新紧邻点所产生的引力向量大小
        old_gravity = None
        for neighbor_index in neighbor_info:
            neighbor = matrix[neighbor_index] # 近邻点坐标
            gravity_vector = neighbor - source # 紧邻点和源点产生的引力向量
            if old_gravity is None:
                old_gravity = gravity_vector
                gravity_vector_size_list.append(vector_module(old_gravity))
            else:
                new_gravity = old_gravity + gravity_vector # 在原来的引力向量基础上增加一个新紧邻点，产生的新引力向量
                gravity_vector_size_list.append(vector_module(new_gravity))
                old_gravity = new_gravity
        
        steps = []
        for i in range(1, len(gravity_vector_size_list)):
            step = gravity_vector_size_list[i] - gravity_vector_size_list[i - 1];
            steps.append(step)
        res.append(steps)
    return res



if __name__ == "__main__":
    filename = "data/data0.txt"
    points =np.loadtxt(filename, delimiter=',')  # type: np.ndarray
    res = get_step_increment_of_gravity(points)
    pprint(res)
    result_file = os.path.splitext(filename)[0]
    print("finished! writen to {:s}".format(result_file))

